import pandas as pd
import matplotlib.pyplot as plt

# Load the data from Excel
file_path = 'RMSD_average_all_residues_f513.xlsx'
df = pd.read_excel(file_path, sheet_name='Sheet1')

# Clean up the dataframe by removing the first row and setting proper column names
df_cleaned = df.iloc[1:].copy()
df_cleaned.columns = ['Frame', 'R375K', 'R375L', 'RW_LF', 'W376A', 'WT']
df_cleaned['Frame'] = df_cleaned['Frame'].astype(int)
df_cleaned[['R375K', 'R375L', 'RW_LF', 'W376A', 'WT']] = df_cleaned[['R375K', 'R375L', 'RW_LF', 'W376A', 'WT']].astype(float)

# Define custom colors
colors = {
    'R375K': 'darkblue',
    'R375L': 'darkgreen',
    'RW_LF': 'lightsteelblue',
    'W376A': 'orange',  # grayscale color for W376A
    'WT': 'black'
}

# Plotting the data with custom colors
plt.figure(figsize=(10, 6))
for column in ['R375K', 'R375L', 'RW_LF', 'W376A', 'WT']:
    plt.plot(df_cleaned['Frame'], df_cleaned[column], label=column, color=colors[column])

# Set x-axis and y-axis limits
plt.xlim(50, 513)
plt.ylim(0, 20)

# Calculate and plot the average RMSD over all frames
average_rmsd = df_cleaned[['R375K', 'R375L', 'RW_LF', 'W376A', 'WT']].mean()

# Plot average RMSD as dotted lines for each mutant with custom colors
#for column in ['R375K', 'R375L', 'W376A', 'WT']:
    #plt.axhline(y=average_rmsd[column], linestyle='--', label=f'Avg {column}', color=colors[column], alpha=0.7)

# Adding labels, title, and legend
plt.xlabel('Frame', labelpad=20, fontsize=20)
plt.ylabel('Average RMSD, Å', labelpad=20, fontsize=20)
plt.title('Average RMSD Å vs Frame for Different Mutants with Overall Average', pad=20, fontsize=18)
plt.legend(loc="upper left", fontsize=12)
plt.xticks(fontsize=16)
plt.yticks([0, 5, 10, 15, 20], fontsize=16)

# Save the plot to a file
plt.savefig('Average_RMSD_for_513_frames_NEW.svg', dpi=500)

# Show the plot
plt.show()

# Print the average RMSD values
print("Average RMSD for each mutant:")
print(average_rmsd)
